from maix import camera , display , image , pinmap , gpio , app,time ,nn
import math , cv2
import numpy as np



detector = nn.YOLOv5(model="/root/models/chess/model_224463.mud", dual_buff = True)
cam = camera.Camera(detector.input_width(), detector.input_height(), detector.input_format(),fps=60)
disp = display.Display()
'''创建按键对象'''
pinmap.set_pin_function("A28","GPIOA28")
pinmap.set_pin_function("A29","GPIOA29")
pinmap.set_pin_function("A19","GPIOA19")
pinmap.set_pin_function("A18","GPIOA18")
key1 = gpio.GPIO("GPIOA19",gpio.Mode.IN)
key2 = gpio.GPIO("GPIOA18",gpio.Mode.IN)
key3 = gpio.GPIO("GPIOA29",gpio.Mode.IN)
key4 = gpio.GPIO("GPIOA28",gpio.Mode.IN)
pinmap.set_pin_function("B3", "GPIOB3")
led = gpio.GPIO("GPIOB3", gpio.Mode.OUT)



run_task = 1  # 1:识别棋盘   2:确定棋盘
#棋盘的四个角点
need_blobs = []
#棋盘九宫格判断字典
chess_list = {  
                1:0  ,  2:0  ,  3:0 , 
                4:0  ,  5:0  ,  6:0 , 
                7:0  ,  8:0  ,  9:0
             }
points = [] # 九宫格十六个坐标点
region = {} # 九宫格数字区域

def check_key():
    global run_task
    if not key1.value():
        print("按键一被按下")
        run_task = 1

    elif not key2.value():
        print("按键二被按下")
        run_task = 2
        
    elif not key3.value():
        print("按键三被按下")
        run_task = 3
        #red_laser.reset_task()
        
    elif not key4.value():
        print("按键四被按下")
        run_task = 4

def get_contour(flag=0):
    points = []

    cv_img_raw = image.image2cv(img) # 转cv
    cv_img = cv2.cvtColor(cv_img_raw,cv2.COLOR_BGR2GRAY) # 转灰度
    #cv_img = cv2.blur(cv_img,(3,3))
    #cv_img = cv2.bilateralFilter(cv_img,9,10,10)     # 双边滤波
    #cv_img = cv2.morphologyEx(cv_img,cv2.MORPH_CLOSE,kernel) # 闭运算
    cv_img = cv2.Canny(cv_img ,200,300) # 边缘检测
    #cv_img = cv2.morphologyEx(cv_img, cv2.MORPH_CLOSE, np.ones((7,7), np.uint8))  # 闭合大裂缝
    cv_img = cv2.dilate(cv_img, np.ones((5,5), np.uint8), iterations=1)          # 加粗边缘
    if flag:
        img_show = image.cv2image(cv_img)
        disp.show(img_show)
        print(111)
    else:
        contours , _ = cv2.findContours(cv_img,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE) # 找轮廓
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 2000:
                continue
            epsilon  = 0.1 * cv2.arcLength(contour,True)  # 多边形逼近
            approx = cv2.approxPolyDP(contour,epsilon,True)
            if len(approx) == 4:
                cv2.drawContours(cv_img_raw,[approx],0,(0,255,0),2)
                quad_points = []
                for point in approx:
                    x , y = point.ravel()
                    quad_points.append([x, y])  # 收集当前四边形的点
                    #cv2.circle(cv_img_raw , ( x , y ) , 5 , ( 255 , 0 , 0 ) , 1)
                
                points = quad_points
        img_show = image.cv2image(cv_img_raw)

    return img_show , points

def generate_16_points(corners):
    """
    根据任意顺序的四个角点生成16个网格点
    参数：corners - 四个点的列表 [(x1,y1), (x2,y2), (x3,y3), (x4,y4)]
    返回：16个网格点，顺序为行优先（左上→右下）
    """
    # 1. 确定物理位置上的角点顺序 (基于坐标值)
    def order_corners_by_coords(pts):
        # 找到最小y值（顶部位置）
        min_y = min(p[1] for p in pts)
        max_y = max(p[1] for p in pts)
        
        # 分离上下边界的点
        top_pts = [p for p in pts if p[1] < (min_y + max_y)/2]
        bottom_pts = [p for p in pts if p[1] >= (min_y + max_y)/2]
        
        # 在顶部点中找到最左和最右的点
        top_pts_sorted = sorted(top_pts, key=lambda p: p[0])
        top_left = top_pts_sorted[0]
        top_right = top_pts_sorted[-1]
        
        # 在底部点中找到最左和最右的点
        bottom_pts_sorted = sorted(bottom_pts, key=lambda p: p[0])
        bottom_left = bottom_pts_sorted[0]
        bottom_right = bottom_pts_sorted[-1]
        
        return [top_left, top_right, bottom_right, bottom_left]
    
    # 2. 获取按物理位置排序的角点 [左上, 右上, 右下, 左下]
    tl, tr, br, bl = order_corners_by_coords(corners)
    
    # 3. 生成16个网格点（双线性插值）
    points = []
    for v in [0, 1/3, 2/3, 1]:  # 垂直方向
        for u in [0, 1/3, 2/3, 1]:  # 水平方向
            # 顶部边缘插值
            top_x = tl[0] + u * (tr[0] - tl[0])
            top_y = tl[1] + u * (tr[1] - tl[1])
            
            # 底部边缘插值
            bottom_x = bl[0] + u * (br[0] - bl[0])
            bottom_y = bl[1] + u * (br[1] - bl[1])
            
            # 垂直插值
            x = top_x + v * (bottom_x - top_x)
            y = top_y + v * (bottom_y - top_y)
            points.append((int(x), int(y)))
    
    return points

def assign_fixed_numbers(grid_points):
    """
    为九宫格分配固定物理位置序号
    参数：grid_points - 16个网格点（由generate_16_points生成）
    返回：9个区域及其固定序号的字典
    """
    # 1. 计算9个区域的中心点
    regions = {}
    for i in range(9):
        row = i // 3
        col = i % 3
        
        # 获取网格点索引
        top_left_idx = row*4 + col
        top_right_idx = row*4 + col + 1
        bottom_left_idx = (row+1)*4 + col
        bottom_right_idx = (row+1)*4 + col + 1
        
        # 计算区域中心
        center_x = sum([
            grid_points[top_left_idx][0],
            grid_points[top_right_idx][0],
            grid_points[bottom_left_idx][0],
            grid_points[bottom_right_idx][0]
        ]) // 4
        
        center_y = sum([
            grid_points[top_left_idx][1],
            grid_points[top_right_idx][1],
            grid_points[bottom_left_idx][1],
            grid_points[bottom_right_idx][1]
        ]) // 4
        
        # 固定位置分配（不受旋转影响）
        positions = {
            0: 1, 1: 2, 2: 3,
            3: 4, 4: 5, 5: 6,
            6: 7, 7: 8, 8: 9
        }
        
        regions[i] = {
            "center": (center_x, center_y),
            "number": positions[i]
        }
    
    return regions

def draw_chess_board(grid_points,region_data):
    global need_blobs
    if not need_blobs or len(need_blobs) != 4:
        return
    
    # # 1. 生成16个网格点
    # grid_points = generate_16_points(need_blobs)
    
    # # 2. 获取9个区域的固定序号信息
    # region_data = assign_fixed_numbers(grid_points)
    
    # 3. 绘制九宫格网格线（蓝色）
    for i in range(4):
        # 水平线
        img.draw_line(grid_points[i*4][0], grid_points[i*4][1],
                     grid_points[i*4+3][0], grid_points[i*4+3][1],
                     color=image.COLOR_BLUE, thickness=2)
        
        # 垂直线
        img.draw_line(grid_points[i][0], grid_points[i][1],
                     grid_points[12+i][0], grid_points[12+i][1],
                     color=image.COLOR_BLUE, thickness=2)
    
    # 4. 绘制固定序号（红色）
    for i in range(9):
        center_x, center_y = region_data[i]["center"]
        number = region_data[i]["number"]  # 提取数字部分
        img.draw_string(center_x-5, center_y-5, f"{number}", 
                      color=image.COLOR_RED, scale=1.5)
        #img.draw_circle(center_x, center_y, 3, color=image.COLOR_GREEN, thickness=-1)
    
def get_position_in_grid(point, grid_points):
    """
    判断坐标点位于九宫格的哪个区域
    参数:
        point: 待判断的坐标点 (x, y)
        grid_points: 16个网格点列表（由generate_16_points生成）
    返回:
        如果点在九宫格内：返回区域编号 (1-9)
        如果点在九宫格外：返回 None
    """
    if not point:
        return

    # 1. 检查点是否在九宫格边界内
    min_x = min(p[0] for p in grid_points)
    max_x = max(p[0] for p in grid_points)
    min_y = min(p[1] for p in grid_points)
    max_y = max(p[1] for p in grid_points)
    
    px, py = point
    if not (min_x <= px <= max_x and min_y <= py <= max_y):
        return None  # 点在九宫格外
    
    # 2. 遍历9个区域判断点所在位置
    for i in range(9):
        row = i // 3
        col = i % 3
        
        # 获取区域四个角点
        top_left = grid_points[row*4 + col]
        top_right = grid_points[row*4 + col+1]
        bottom_left = grid_points[(row+1)*4 + col]
        bottom_right = grid_points[(row+1)*4 + col+1]
        
        # 3. 判断点是否在当前区域内（使用射线法）
        if is_point_in_quad(point, top_left, top_right, bottom_right, bottom_left):
            return i + 1  # 返回区域编号（1-9）
    
    return None  # 点在网格线上或边界外

def is_point_in_quad(p, a, b, c, d):
    """
    判断点p是否在四边形abcd内（顺时针顺序）
    使用射线法（从点向右水平发射射线）
    """
    def cross_product(v1, v2):
        """计算向量叉积"""
        return v1[0]*v2[1] - v1[1]*v2[0]
    
    def vector(from_point, to_point):
        """计算向量"""
        return (to_point[0]-from_point[0], to_point[1]-from_point[1])
    
    # 检查四条边
    ab = vector(a, b)
    ap = vector(a, p)
    bc = vector(b, c)
    bp = vector(b, p)
    cd = vector(c, d)
    cp = vector(c, p)
    da = vector(d, a)
    dp = vector(d, p)
    
    # 计算叉积符号
    cross1 = cross_product(ab, ap)
    cross2 = cross_product(bc, bp)
    cross3 = cross_product(cd, cp)
    cross4 = cross_product(da, dp)
    
    # 如果所有叉积符号相同（都在边的同侧），则点在四边形内
    if (cross1 >= 0 and cross2 >= 0 and cross3 >= 0 and cross4 >= 0) or \
       (cross1 <= 0 and cross2 <= 0 and cross3 <= 0 and cross4 <= 0):
        return True
    
    return False



led.value(1)
while not app.need_exit():
    img = cam.read()
    #img.binary(chess_board_threshold)
    check_key()
    target_point = []

    objs = detector.detect(img, conf_th = 0.3, iou_th = 0.2)
    for obj in objs:
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color = image.COLOR_RED)
        msg = f'{detector.labels[obj.class_id]}: {obj.score:.2f}'
        img.draw_string(obj.x, obj.y, msg, color = image.COLOR_RED)
        
        target_point = [int(obj.x + (obj.w/2)) , int(obj.y + (obj.h/2)) ]
        print(obj.x , obj.y , obj.w , obj.h , target_point)
        img.draw_keypoints(target_point,color=image.COLOR_GREEN)
        

    if run_task == 1:
        img , need_blobs = get_contour()
    elif run_task == 2:
        if need_blobs:
            points = generate_16_points(need_blobs)
            region = assign_fixed_numbers(points)
            draw_chess_board(points,region)

            
            index = get_position_in_grid(target_point,points)
            if index:
                print(index)
                chess_list[index] = 1
                print(chess_list)
    elif run_task == 3:
        draw_chess_board(points,region)
        img.draw_keypoints([50,50],color=image.COLOR_GREEN)
        

        
        


    print(time.fps())
    
    disp.show(img)